# 构建和训练deep learning模型
import torch
import torch.nn as nn # 神经网络层
from torch.utils.data import Dataset,DataLoader # 用来对数据进行封装

# 对数据操作
import numpy as np # NumPy科学计算
import csv
import os

# 绘图相关库
from matplotlib.pyplot import figure # 用于创建图形对象
import matplotlib.pyplot as plt # Matplotlib绘图库，pyplot提供了类似MATLAB的绘图接口

# 


def get_device():
    if torch.cuda.is_available():
        return 'cuda'
    else:
        return 'cpu'
  
# train 和 dev loss图像绘制
# 参数说明：train or dev，自定义的title
def train_dev_learning_loss_plt(loss_record,title=''):
  # 训练集处理
    # 训练集步数
    total_steps = len(loss_record['train'])
    x_1 = range(total_steps)
    # 设置图像大小
    figure(figsize=(6,4))
    # 绘制训练集损失折线图
    # 横坐标，训练损失，线条为红色，为曲线设置标签
    plt.plot(x_1,loss_record['train'],c="tab:red",label='train') 
  # 开发集处理(节省计算资源)
    if len(loss_record['rev']) != 0:
        x_2 = x_1[::len(loss_record['train'])//len(loss_record['dev'])]
        plt.plot(x_2,loss_record['dev'],c='tab:cyan',label='dev')
  # 图像设置参数
    plt.ylim(0.0,20.0) # 强制Y轴显示范围为 0.0 到 20.0
    plt.xlabel('Training steps')
    plt.ylabel('MSE loss')
    plt.title('Learning curve of {}'.format(title))
  # 显示图像
    plt.legend() # 显示train和dev折线图
    plt.show() # 弹出窗口显示图像

# 可视化深度神经网络（DNN）模型
# 绘制模型预测值与真实值的散点图，用于直观评估回归模型的准确性
# 参数说明：开发集，训练模型，cpu or gpu，横坐标范围，模型的预测值，实际目标值
def plot_pred_true(devset,model,device,lim=35.,preds=None,targets=None):
    if preds is None or targets is None:
        model.eval() # 关闭随机性
        preds,targets = [],[]
        for x,y in devset: # x是输入特征，y是实际目标值
            x,y=x.to(device),y.to(device)
            with torch.no_grad(): # 我们不需要算梯度（改参数的话，误差会怎么变）
                pred = model(x)   # 我们将一个训练好的模型，输入数据，让它猜一下结果
                # pred.detach().cpu() # .detach()就是断开连接
                preds.append(pred.detach().cpu()) 
                targets.append(y.detach().cpu())
            preds = torch.cat(preds,dim=0).numpy()
            targets = torch.cat(targets,dim=0).numpy()
        
        figure(figsize=(5,5))
        plt.scatter(targets,preds,c='r',alpha=0.5) # 画一个散点图，如果准确，就应该落在对角线上
        plt.plot([-0.2,lim],[-0.2,lim],c='b')
        plt.xlim(-0.2,lim)
        plt.ylim(-0.2,lim)

        plt.xlabel('ground truth value')           # x轴标签
        plt.ylabel('predicted value')              # y轴标签
        plt.title('Ground Truth v.s. Prediction')  # 图的标题
        plt.show()

# 数据集加载及预处理
class COVID19Dataset(Dataset):
    def __init__(self,path,mode='train',target_only=False):
        self.mode = mode

        # 读取数据
        with open(path,'r') as fp:
            data = list(csv.reader(fp))
            data = np.array(data[1:])[:,1:].astype(float)
        
        # 选择使用哪些特征
        if not target_only:
            feats = list(range(93)) # 使用前93个特征
        else:
            feats = [40, 41, 42, 43, 57, 58, 59, 60, 61, 75, 76, 77, 78, 79] # 网络数据

        # 根据mode分情况处理
        if mode == 'test':
            data = data[:,feats] # 只保留特征列
            self.data = torch.FloatTensor(data) # 转成 pytorch 张量
        else:
            # train/dev
            data = data[:,feats]
            target = data[:,-1]

            self.mean = torch.FloatTensor(data).mean(dim=0,keepdim=True) # 均值
            self.std = torch.FloatTensor(data).std(dim=0,keepdim=True) # 标准差
 
            # 总数据按顺序编号，每5个中取1个作为验证集，其余4个做训练集
            if mode == 'train':
                indices = [i for i in range(len(data)) if i % 5 != 0]
            elif mode == 'dev':
                indices = [i for i in range(len(data)) if i % 5 == 0]

            self.data = torch.FloatTensor(data[indices])
            self.target = torch.FloatTensor(data[indices])

        self.dim = self.data.shape[1] # 获取数据集self.data的列数，也就是特征的数量

        print('Finished reading the {} set of COVID19 Dataset ({} samples found, each dim = {})'.format(mode, len(self.data), self.dim))


    # 根据一个编号，取出一条数据样本
    def __getitem__(self, index):
        if self.mode in ['train','dev']:
            return self.data[index],self.target[index]
        else:
            return self.data[index]
        
    def __len__(self):
        return len(self.data)
    
    def normalization(self,mean = None,std = None):
        if self.mode == 'train' or self.mode == 'dev':
            mean = self.mean
            std = self.std
            self.data = (self.data - mean) / std
        else:
            self.data = (self.data - mean) / std

        return mean,std
    
# 构造数据加载集，用于训练。验证或测试机器学习模型
def prep_dataloader(path,mode,batch_size,n_jobs,target_only=False,mean=False,std=None):
    dataset = COVID19Dataset(path,mode=mode,target_only=target_only)
    mean,std=dataset.normalization(mean,std)
    dataloader = DataLoader(dataset, batch_size,shuffle=(mode == 'train'), drop_last=False,num_workers=n_jobs, pin_memory=True) 
    return dataloader,mean,std

# 构建前馈神经网络（Feedforward Neural Network）模型
# 简单的全连接神经网络
class NeuralNet(nn.Module):
    # 网络结构
    def __init__(self,input_dim):
        super(NeuralNet,self).__init__()
        self.net = nn.Sequential(
            nn.Linear(input_dim,64), # 输入层 -> 隐藏层1
            nn.ReLU(),
            nn.Linear(64,16), # 隐藏层1 -> 隐藏层2
            nn.ReLU(),
            nn.Linear(16,8), # 隐藏层2 -> 隐藏层3
            nn.ReLU(),
            nn.Linear(8,4), # 隐藏层3 -> 隐藏层4
            nn.ReLU(),
            nn.Linear(4,1) # 隐藏层4 -> 输出层
        )
        # 损失函数
        # MSE用于回归任务，计算预测值和真实值的平方误差均值
        self.criterion = nn.MSELoss(reduction='mean') # 均方误差损失
        
    #








